/* SPDX-License-Identifier: GPL-2.0 */
/*
 * 一个简单的调度器，用于演示用户态和内核态BPF程序之间的通信。
 *
 * 这个调度器基于scx_simple实现，添加了用户态和内核态之间的通信功能。
 * 用户态程序可以写入一个消息，内核态BPF程序会读取并打印这个消息。
 * 同时支持传递 top10 的 LLC 缓存未命中率线程信息。
 *
 * Copyright (c) 2024
 */
#include <scx/common.bpf.h>
#include "scx_hello.h"

char _license[] SEC("license") = "GPL";

const volatile bool fifo_sched;

static u64 vtime_now;
UEI_DEFINE(uei);

/*
 * 共享DSQ，与scx_simple相同
 */
#define SHARED_DSQ 0

/*
 * 用于存储用户态传递的消息
 */
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(key_size, sizeof(u32));
    __uint(value_size, sizeof(struct scx_hello_message));
    __uint(max_entries, 1);
} hello_message SEC(".maps");

/*
 * 用于存储用户态传递的 LLC 缓存未命中率线程信息
 */
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(key_size, sizeof(u32));
    __uint(value_size, sizeof(struct scx_hello_llc_info));
    __uint(max_entries, 1);
} llc_miss_info SEC(".maps");

/*
 * 统计信息，与scx_simple相同
 */
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __uint(key_size, sizeof(u32));
    __uint(value_size, sizeof(u64));
    __uint(max_entries, 2);         /* [local, global] */
} stats SEC(".maps");

static void stat_inc(u32 idx)
{
    u64 *cnt_p = bpf_map_lookup_elem(&stats, &idx);
    if (cnt_p)
        (*cnt_p)++;
}

/*
 * 检查并打印用户态传递的消息
 */
static void check_and_print_message(void)
{
    u32 key = 0;
    struct scx_hello_message *msg;
    
    msg = bpf_map_lookup_elem(&hello_message, &key);
    if (!msg)
        return;
    
    if (msg->new_message) {
        bpf_printk("Hello from user space: %s\n", msg->message);
        msg->new_message = false;
    }
}

/*
 * 检查并打印用户态传递的 LLC 缓存未命中率线程信息
 */
static void check_and_print_llc_miss_info(void)
{
    u32 key = 0;
    struct scx_hello_llc_info *info;
    int i;
    
    info = bpf_map_lookup_elem(&llc_miss_info, &key);
    if (!info)
        return;
    
    if (info->new_data) {
        bpf_printk("LLC Cache Miss Ratio Top Threads (count: %d):\n", info->count);
        
        /* 使用循环展开来避免 BPF 验证器的限制 */
        if (info->count > 0) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[0].tid, info->threads[0].miss_ratio);
        if (info->count > 1) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[1].tid, info->threads[1].miss_ratio);
        if (info->count > 2) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[2].tid, info->threads[2].miss_ratio);
        if (info->count > 3) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[3].tid, info->threads[3].miss_ratio);
        if (info->count > 4) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[4].tid, info->threads[4].miss_ratio);
        if (info->count > 5) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[5].tid, info->threads[5].miss_ratio);
        if (info->count > 6) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[6].tid, info->threads[6].miss_ratio);
        if (info->count > 7) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[7].tid, info->threads[7].miss_ratio);
        if (info->count > 8) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[8].tid, info->threads[8].miss_ratio);
        if (info->count > 9) bpf_printk("  TID: %d, Miss Ratio: %.2f%%\n", info->threads[9].tid, info->threads[9].miss_ratio);
        
        info->new_data = false;
    }
}

s32 BPF_STRUCT_OPS(hello_select_cpu, struct task_struct *p, s32 prev_cpu, u64 wake_flags)
{
    bool is_idle = false;
    s32 cpu;

    cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, wake_flags, &is_idle);
    if (is_idle) {
        stat_inc(0);    /* count local queueing */
        scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
    }

    return cpu;
}

void BPF_STRUCT_OPS(hello_enqueue, struct task_struct *p, u64 enq_flags)
{
    stat_inc(1);    /* count global queueing */

    /* 检查并打印用户态传递的消息 */
    check_and_print_message();
    
    /* 检查并打印用户态传递的 LLC 缓存未命中率线程信息 */
    check_and_print_llc_miss_info();

    if (fifo_sched) {
        scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, enq_flags);
    } else {
        u64 vtime = p->scx.dsq_vtime;

        /*
         * Limit the amount of budget that an idling task can accumulate
         * to one slice.
         */
        if (time_before(vtime, vtime_now - SCX_SLICE_DFL))
            vtime = vtime_now - SCX_SLICE_DFL;

        scx_bpf_dsq_insert_vtime(p, SHARED_DSQ, SCX_SLICE_DFL, vtime,
                     enq_flags);
    }
}

void BPF_STRUCT_OPS(hello_dispatch, s32 cpu, struct task_struct *prev)
{
    scx_bpf_dsq_move_to_local(SHARED_DSQ);
}

void BPF_STRUCT_OPS(hello_running, struct task_struct *p)
{
    if (fifo_sched)
        return;

    /*
     * Global vtime always progresses forward as tasks start executing. The
     * test and update can be performed concurrently from multiple CPUs and
     * thus racy. Any error should be contained and temporary. Let's just
     * live with it.
     */
    if (time_before(vtime_now, p->scx.dsq_vtime))
        vtime_now = p->scx.dsq_vtime;
}

void BPF_STRUCT_OPS(hello_stopping, struct task_struct *p, bool runnable)
{
    if (fifo_sched)
        return;

    /*
     * Scale the execution time by the inverse of the weight and charge.
     *
     * Note that the default yield implementation yields by setting
     * @p->scx.slice to zero and the following would treat the yielding task
     * as if it has consumed all its slice. If this penalizes yielding tasks
     * too much, determine the execution time by taking explicit timestamps
     * instead of depending on @p->scx.slice.
     */
    p->scx.dsq_vtime += (SCX_SLICE_DFL - p->scx.slice) * 100 / p->scx.weight;
}

void BPF_STRUCT_OPS(hello_enable, struct task_struct *p)
{
    p->scx.dsq_vtime = vtime_now;
}

s32 BPF_STRUCT_OPS_SLEEPABLE(hello_init)
{
    /* 初始化消息映射 */
    u32 key = 0;
    struct scx_hello_message msg = {0};
    struct scx_hello_llc_info llc_info = {0};
    
    bpf_map_update_elem(&hello_message, &key, &msg, 0);
    bpf_map_update_elem(&llc_miss_info, &key, &llc_info, 0);
    
    return scx_bpf_create_dsq(SHARED_DSQ, -1);
}

void BPF_STRUCT_OPS(hello_exit, struct scx_exit_info *ei)
{
    UEI_RECORD(uei, ei);
}

SCX_OPS_DEFINE(hello_ops,
           .select_cpu        = (void *)hello_select_cpu,
           .enqueue           = (void *)hello_enqueue,
           .dispatch          = (void *)hello_dispatch,
           .running           = (void *)hello_running,
           .stopping          = (void *)hello_stopping,
           .enable            = (void *)hello_enable,
           .init              = (void *)hello_init,
           .exit              = (void *)hello_exit,
           .name              = "hello");